#coding=utf-8
from xml.dom.minidom import Document    #使用这个工具来读写xml
import random
class URDF_GEN:

    def __init__(self,name='unamed'):
        '''配置机器人'''
        objNums = 0  # 每新建一个元件id就+1
        self.doc=Document()    #机器人的描述性文件
        self.robot=self.doc.createElement("robot")
        self.robot.setAttribute('name',name)
        self.tree={}
    def Add_Cylinder(self,name='',
                     len=0.0,
                     radius=0.0,
                     father='',
                     origin_xyz='0.0 0.0 0.0',
                     origin_rpy='0.0 0.0 0.0',
                     rgba="1 1 1 1"
                     ):
        '''添加一个圆柱形'''
        if father=='': father=self.robot
        subLink=self.doc.createElement('link')
        subLink.setAttribute('name', name)

        #visual部分
        visual = self.doc.createElement('visual')

        #geometry部分
        geom = self.doc.createElement('geometry')
        cyln = self.doc.createElement('cylinder')
        cyln.setAttribute('length',str(len))
        cyln.setAttribute('radius', str(radius))
        geom.appendChild(cyln)
        visual.appendChild(geom)
        #material部分
        material = self.doc.createElement('material')
        color = self.doc.createElement('color')
        material.setAttribute('name', 'color_'+str(int(random.random()*10000)))
        color.setAttribute('rgba', rgba)
        material.appendChild(color)
        visual.appendChild(material)
        #origin部分
        origin = self.doc.createElement('origin')
        origin.setAttribute('rpy', origin_rpy)
        origin.setAttribute('xyz', origin_xyz)
        visual.appendChild(origin)
        #collision
        geom = self.doc.createElement('geometry')
        cyln = self.doc.createElement('cylinder')
        cyln.setAttribute('length', str(len))
        cyln.setAttribute('radius', str(radius))
        geom.appendChild(cyln)
        origin = self.doc.createElement('origin')
        origin.setAttribute('rpy', origin_rpy)
        origin.setAttribute('xyz', origin_xyz)
        coll = self.doc.createElement('collision')
        coll.appendChild(geom)
        coll.appendChild(origin)

        #添加
        subLink.appendChild(visual)
        subLink.appendChild(coll)
        father.appendChild(subLink)
    def Add_Cube(self,name='',
                     sSize='0 0 0',
                     father='',
                     origin_xyz='0.0 0.0 0.0',
                     origin_rpy='0.0 0.0 0.0',
                     rgba="1 1 1 1"
                     ):
        '''添加一个长方形'''
        if father == '': father = self.robot
        subLink = self.doc.createElement('link')
        subLink.setAttribute('name', name)

        # visual部分
        visual = self.doc.createElement('visual')

        # geometry部分
        geom = self.doc.createElement('geometry')
        bbx = self.doc.createElement('box')
        bbx.setAttribute('size', sSize)
        geom.appendChild(bbx)
        visual.appendChild(geom)
        # material部分
        material = self.doc.createElement('material')
        color = self.doc.createElement('color')
        material.setAttribute('name', 'color_' + str(int(random.random() * 10000)))
        color.setAttribute('rgba', rgba)
        material.appendChild(color)
        visual.appendChild(material)
        # origin部分
        origin = self.doc.createElement('origin')
        origin.setAttribute('rpy', origin_rpy)
        origin.setAttribute('xyz', origin_xyz)
        visual.appendChild(origin)
        # collision
        geom = self.doc.createElement('geometry')
        bbx = self.doc.createElement('box')
        bbx.setAttribute('size', sSize)
        geom.appendChild(bbx)
        origin = self.doc.createElement('origin')
        origin.setAttribute('rpy', origin_rpy)
        origin.setAttribute('xyz', origin_xyz)
        coll = self.doc.createElement('collision')
        coll.appendChild(geom)
        coll.appendChild(origin)

        # 添加
        subLink.appendChild(visual)
        subLink.appendChild(coll)
        father.appendChild(subLink)
    def Add_Mesh(self,name='',
                     meshFile='',
                     father='',
                     origin_xyz='0.0 0.0 0.0',
                     origin_rpy='0.0 0.0 0.0',
                     rgba="1 1 1 1"
                     ):
        '''添加一个圆柱形'''
        '''添加一个球形'''
        if father == '': father = self.robot
        subLink = self.doc.createElement('link')
        subLink.setAttribute('name', name)

        # visual部分
        visual = self.doc.createElement('visual')

        # geometry部分
        geom = self.doc.createElement('geometry')
        mesh = self.doc.createElement('mesh')
        mesh.setAttribute('filename', meshFile)
        geom.appendChild(mesh)
        visual.appendChild(geom)
        # material部分
        material = self.doc.createElement('material')
        color = self.doc.createElement('color')
        material.setAttribute('name', 'color_' + str(int(random.random() * 10000)))
        color.setAttribute('rgba', rgba)
        material.appendChild(color)
        visual.appendChild(material)
        # origin部分
        origin = self.doc.createElement('origin')
        origin.setAttribute('rpy', origin_rpy)
        origin.setAttribute('xyz', origin_xyz)
        visual.appendChild(origin)
        # collision
        geom = self.doc.createElement('geometry')
        mesh = self.doc.createElement('mesh')
        mesh.setAttribute('filename', meshFile)
        geom.appendChild(mesh)
        origin = self.doc.createElement('origin')
        origin.setAttribute('rpy', origin_rpy)
        origin.setAttribute('xyz', origin_xyz)
        coll = self.doc.createElement('collision')
        coll.appendChild(geom)
        coll.appendChild(origin)

        # 添加
        subLink.appendChild(visual)
        subLink.appendChild(coll)
        father.appendChild(subLink)
    def Add_Cube(self,name='',
                     sSize='0 0 0',
                     father='',
                     origin_xyz='0.0 0.0 0.0',
                     origin_rpy='0.0 0.0 0.0',
                     rgba="1 1 1 1"
                     ):
        '''添加一个长方形'''
        if father == '': father = self.robot
        subLink = self.doc.createElement('link')
        subLink.setAttribute('name', name)

        # visual部分
        visual = self.doc.createElement('visual')

        # geometry部分
        geom = self.doc.createElement('geometry')
        bbx = self.doc.createElement('box')
        bbx.setAttribute('size', sSize)
        geom.appendChild(bbx)
        visual.appendChild(geom)
        # material部分
        material = self.doc.createElement('material')
        color = self.doc.createElement('color')
        material.setAttribute('name', 'color_' + str(int(random.random() * 10000)))
        color.setAttribute('rgba', rgba)
        material.appendChild(color)
        visual.appendChild(material)
        # origin部分
        origin = self.doc.createElement('origin')
        origin.setAttribute('rpy', origin_rpy)
        origin.setAttribute('xyz', origin_xyz)
        visual.appendChild(origin)
        # collision
        geom = self.doc.createElement('geometry')
        bbx = self.doc.createElement('box')
        bbx.setAttribute('size', sSize)
        geom.appendChild(bbx)
        origin = self.doc.createElement('origin')
        origin.setAttribute('rpy', origin_rpy)
        origin.setAttribute('xyz', origin_xyz)
        coll = self.doc.createElement('collision')
        coll.appendChild(geom)
        coll.appendChild(origin)

        # 添加
        subLink.appendChild(visual)
        subLink.appendChild(coll)
        father.appendChild(subLink)
    def Add_Joint(self,name='unamed_J',
                  type='revolute',
                  parent_link='',
                  child_link='',
                  origin_rpy='0 0 0',
                  origin_xyz='0 0 0',
                  axis_xyz='0 0 1',
                  vel=0.5,
                  max=0.785,
                  min=-0.785,
                  effort=1000.0,
                  father=''
                  ):
        '''
            添加一个关节
            :param name:
            :param type:
            :param parent_link:
            :param child_link:
            :param origin_rpy:
            :param origin_xyz:
            :param vel:
            :param max:
            :param min:
        :return:
        '''
        if father == '': father = self.robot
        subJoint = self.doc.createElement('joint')
        subJoint.setAttribute('name', name)
        subJoint.setAttribute('type', type)
        #axis
        axis = self.doc.createElement('axis')
        axis.setAttribute('xyz',axis_xyz)
        subJoint.appendChild(axis)
        # origin
        ori = self.doc.createElement('origin')
        ori.setAttribute('rpy', origin_rpy)
        ori.setAttribute('xyz', origin_xyz)
        subJoint.appendChild(ori)
        #parent/child
        par = self.doc.createElement('parent')
        chi= self.doc.createElement('child')
        par.setAttribute('link',parent_link)
        chi.setAttribute('link',child_link)
        subJoint.appendChild(par)
        subJoint.appendChild(chi)
        if type=='fixed':
            self.robot.appendChild(subJoint)
            return
        #limit
        lim = self.doc.createElement('limit')
        lim.setAttribute('lower', str(min))
        lim.setAttribute('upper', str(max))
        lim.setAttribute('velocity', str(vel))
        lim.setAttribute('effort', str(effort))
        subJoint.appendChild(lim)
        self.robot.appendChild(subJoint)

    def Save(self,fileName='urdf.urdf'):
        self.doc.appendChild(self.robot)
        with open(fileName,'w',encoding='utf-8') as f:
            self.doc.writexml(f,indent='',addindent=' ',newl='\n',encoding='utf-8')
def main():
    u=URDF_GEN('robot')  #生成一台机器人

    #添加Link
    u.Add_Cube(name='base_link',sSize='0.5 0.4 0.7',origin_xyz='0 0 0.35',rgba='0.8 0.8 0.8 1')   #主体
    u.Add_Cube(name='lidar_base',sSize='0.12 0.12 0.25',origin_xyz='0 0 0.125',rgba='0.8 0.8 0.8 1')   #雷达座子
    u.Add_Cylinder(name='lidar_link',len=0.09,radius=0.10,origin_xyz='0 0 0.045',rgba='1.0 0.5 0.0 1')  #雷达
    u.Add_Cylinder(name='cam_base',len=0.15,radius=0.075,origin_xyz='0 0 0.075',rgba='0.8 0.8 0.8 1')  #摄像头座子
    u.Add_Cylinder(name='cam_link',len=0.05,radius=0.05,origin_xyz='0 0 0.025',rgba='1.0 0.0 0.6 1')  #摄像头
    u.Add_Cylinder(name='fl_wheel',len=0.12,radius=0.15,origin_xyz='0 0 0.06',origin_rpy='1.5707 0 0',rgba='0.0 0.0 0.6 1')  #轮
    u.Add_Cylinder(name='fr_wheel',len=0.12,radius=0.15,origin_xyz='0 0 0.06',origin_rpy='-1.5707 0 0',rgba='0.0 0.0 0.6 1')  #轮
    u.Add_Cylinder(name='bl_wheel',len=0.12,radius=0.15,origin_xyz='0 0 0.06',origin_rpy='1.5707 0 0',rgba='0.0 0.0 0.6 1')  #轮
    u.Add_Cylinder(name='br_wheel',len=0.12,radius=0.15,origin_xyz='0 0 0.06',origin_rpy='-1.5707 0 0',rgba='0.0 0.0 0.6 1')  #轮

    #添加Joint
    u.Add_Joint(name='lidar_j1',type='fixed',parent_link='base_link',child_link='lidar_base',origin_xyz='0 0 0.7')
    u.Add_Joint(name='lidar_j2',type='fixed',parent_link='lidar_base',child_link='lidar_link',origin_xyz='0 0 0.25')

    u.Add_Joint(name='cam_j1',type='fixed',parent_link='base_link',child_link='cam_base',origin_xyz='0.17 0 0.7')
    u.Add_Joint(name='cam_j2',type='fixed',parent_link='cam_base',child_link='cam_link',origin_xyz='0 0 0.15',axis_xyz='0 0 1')

    u.Add_Joint(name='jw1',type='fixed',parent_link='base_link',child_link='fl_wheel',axis_xyz='0 1 0',origin_xyz='0.3 0.25 0.15')
    u.Add_Joint(name='jw3',type='fixed',parent_link='base_link',child_link='fr_wheel',axis_xyz='0 -1 0',origin_xyz='0.3 -0.25 0.15')
    u.Add_Joint(name='jw2',type='fixed',parent_link='base_link',child_link='bl_wheel',axis_xyz='0 1 0',origin_xyz='-0.3 0.25 0.15')
    u.Add_Joint(name='jw4',type='fixed',parent_link='base_link',child_link='br_wheel',axis_xyz='0 -1 0',origin_xyz='-0.3 -0.25 0.15')

    u.Save()


# u.Add_Mesh(name='base_link',meshFile='package://simulator_bringup/cfg/small_robo.stl',origin_rpy='1.57 0 1.57',origin_xyz='-0.0875 -0.0925 0',rgba='0.9 0.9 0.9 1')   #底座
# u.Add_Cylinder(name='wheel_L',len=0.02,radius=0.06,origin_xyz='0 0 0.0',rgba='0 0 0 1')
# u.Add_Cylinder(name='wheel_R',len=0.02,radius=0.06,origin_xyz='0 0 0.0',rgba='0 0 0 1')
# u.Add_Cylinder(name='laser',len=0.03,radius=0.04,origin_xyz='0 0 0.02',rgba='0 0.0 1.0 1')
# u.Add_Joint(name='base_laser',type='fixed',parent_link='base_link',child_link='laser',origin_xyz='0.0 0 0.15')
# u.Add_Joint(name='base_l_wheel',type='continuous',parent_link='base_link',child_link='wheel_L',origin_xyz='0 -0.0925 0.04',origin_rpy='1.57 0 0')
# u.Add_Joint(name='base_r_wheel',type='continuous',parent_link='base_link',child_link='wheel_R',origin_xyz='0 0.0925 0.04',origin_rpy='-1.57 0 0')
# u.Save()
